package com.github.common.extend.demo.util.db;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import lombok.Data;
import lombok.experimental.Accessors;
import org.apache.poi.ss.usermodel.Cell;
import org.apache.poi.ss.usermodel.Row;
import org.apache.poi.ss.usermodel.Sheet;
import org.apache.poi.ss.usermodel.Workbook;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;
import java.io.IOException;
import java.io.OutputStream;
import java.sql.*;

/**
 * @author X1993
 * @date 2024/6/29
 * @description
 */
public class JdbcExportUtils {

    /**
     * 根据sql查询的视图生成json
     * @param conn
     * @param timeoutSeconds
     * @param querySql
     */
    public static JSONArray export2Json(Connection conn ,String querySql ,Integer timeoutSeconds) throws SQLException
    {
        Statement stmt = conn.createStatement();
        if (timeoutSeconds != null && timeoutSeconds > 0) {
            stmt.setQueryTimeout(timeoutSeconds);
        }

        ResultSet rs = stmt.executeQuery(querySql);

        JSONArray jsonArray = new JSONArray();

        // 获取列名并写入标题行
        ResultSetMetaData metaData = rs.getMetaData();
        int columnCount = metaData.getColumnCount();
        String[] propertyNames = new String[columnCount];
        for (int i = 1; i <= columnCount; i++) {
            propertyNames[i - 1] = metaData.getColumnName(i);
        }

        // 写入数据行
        while (rs.next()) {
            JSONObject jsonObject = new JSONObject();
            jsonArray.add(jsonObject);
            for (int i = 1; i <= columnCount; i++) {
                jsonObject.put(propertyNames[i - 1], rs.getString(i));
            }
        }

        return jsonArray;
    }

    /**
     * 根据sql查询的视图生成json
     * @param conn
     * @param querySql
     */
    public static JSONArray export2Json(Connection conn ,String querySql) throws SQLException
    {
        return export2Json(conn ,querySql ,null);
    }

    /**
     * 根据sql查询的视图生成excel写入指定流中
     * @param outputStream
     * @param conn
     * @param timeoutSeconds
     * @param sql2SheetParams
     */
    public static void export2Excel(OutputStream outputStream ,Connection conn ,
                                    Integer timeoutSeconds ,Sql2SheetParam... sql2SheetParams)
            throws IOException, SQLException
    {
        Workbook workbook = null;
        try {
            workbook = new XSSFWorkbook();
            Statement stmt = conn.createStatement();
            if (timeoutSeconds != null && timeoutSeconds > 0) {
                stmt.setQueryTimeout(timeoutSeconds);
            }

            for (Sql2SheetParam sql2SheetParam : sql2SheetParams) {
                Sheet sheet = workbook.createSheet(sql2SheetParam.getSheetName());
                resultSet2Sheet(stmt ,sheet , sql2SheetParam.getQuerySql());
            }

            workbook.write(outputStream);
        } finally {
            if (workbook != null) {
                try {
                    workbook.close();
                } catch (IOException e) {
                    throw new IllegalStateException(e);
                }
            }
            try {
                outputStream.close();
            } catch (IOException e) {
                throw new IllegalStateException(e);
            }
        }
    }

    /**
     * 根据sql查询的视图生成excel写入指定流中
     * @param outputStream
     * @param conn
     * @param sql2SheetParams
     */
    public static void export2Excel(OutputStream outputStream ,Connection conn ,Sql2SheetParam... sql2SheetParams)
            throws IOException, SQLException
    {
        export2Excel(outputStream ,conn ,null ,sql2SheetParams);
    }

    private static void resultSet2Sheet(Statement stmt ,Sheet sheet ,String querySql) throws SQLException {
        ResultSet rs = stmt.executeQuery(querySql);
        Row headerRow = sheet.createRow(0); // 创建第一行作为标题行

        // 获取列名并写入标题行
        ResultSetMetaData metaData = rs.getMetaData();
        int columnCount = metaData.getColumnCount();
        for (int i = 1; i <= columnCount; i++) {
            Cell cell = headerRow.createCell(i - 1);
            cell.setCellValue(metaData.getColumnName(i));
        }

        // 写入数据行
        int rowIndex = 1;
        while (rs.next()) {
            Row dataRow = sheet.createRow(rowIndex++);
            for (int i = 1; i <= columnCount; i++) {
                Cell cell = dataRow.createCell(i - 1);
                cell.setCellValue(rs.getString(i));
            }
        }
    }

    @Accessors(chain = true)
    @Data
    public static class Sql2SheetParam {

        private String sheetName;

        private String querySql;

    }

}
